from ai_query import AIQueryManager
from ollama_query import OllamaQueryManager
import json
import traceback
from datetime import datetime, date
from decimal import Decimal
import os
import argparse

class DateTimeEncoder(json.JSONEncoder):
    """
    自定义JSON编码器，支持datetime、date和Decimal类型
    """
    def default(self, obj):
        if isinstance(obj, (datetime, date)):
            return obj.isoformat()
        if isinstance(obj, Decimal):
            return str(obj)
        return super().default(obj)

def save_result_to_file(question: str, result: dict, output_dir: str):
    """
    将查询结果保存到文件
    """
    # 确保输出目录存在
    os.makedirs(output_dir, exist_ok=True)
    
    # 生成文件名（使用时间戳避免重名）
    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    filename = f"query_result_{timestamp}.txt"
    filepath = os.path.join(output_dir, filename)
    
    # 准备要保存的内容
    content = [
        f"问题: {question}",
        "-" * 50,
    ]
    
    if result['success']:
        content.extend([
            "生成的SQL:",
            result.get('sql', '无SQL信息'),
            "\n查询结果:",
            json.dumps(result['results'], ensure_ascii=False, indent=2, cls=DateTimeEncoder)
        ])
    else:
        content.extend([
            f"查询失败: {result['error']}",
            "\n错误详情:",
            result.get('traceback', '无详细信息')
        ])
    
    content.append("-" * 50)
    
    # 写入文件
    with open(filepath, 'w', encoding='utf-8') as f:
        f.write('\n'.join(content))
    
    return filepath

def print_result(question: str, result: dict, output_dir: str = None):
    """
    打印查询结果并保存到文件（如果指定了输出目录）
    """
    print("-" * 50)
    
    if result['success']:
        print("\n查询结果:")
        print(json.dumps(result['results'], ensure_ascii=False, indent=2, cls=DateTimeEncoder))
    else:
        print(f"查询失败: {result['error']}")
        if 'traceback' in result:
            print("\n错误详情:")
            print(result['traceback'])
    
    print("-" * 50)
    
    # 如果指定了输出目录，保存结果到文件
    if output_dir:
        filepath = save_result_to_file(question, result, output_dir)
        print(f"\n结果已保存到文件: {filepath}")

def main():
    try:
        # 解析命令行参数
        parser = argparse.ArgumentParser(description='自然语言转SQL查询工具')
        parser.add_argument('--model', type=str, choices=['deepseek', 'ollama'], default='ollama',
                            help='选择使用的模型服务: deepseek 或 ollama (默认)')
        parser.add_argument('--interactive', action='store_true', 
                            help='交互式选择模型服务')
        args = parser.parse_args()
        
        # 如果指定了交互式模式，让用户选择模型
        if args.interactive:
            print("请选择要使用的模型服务:")
            print("1. DeepSeek API")
            print("2. 本地OLLAMA服务")
            choice = input("请输入选项 (1/2): ").strip()
            if choice == '1':
                args.model = 'deepseek'
            elif choice == '2':
                args.model = 'ollama'
            else:
                print("无效的选项，使用默认模型: ollama")
                args.model = 'ollama'
        
        # 创建输出目录
        output_dir = "query_results"
        
        # 根据选择的模型服务创建相应的查询管理器实例
        if args.model == 'deepseek':
            print("正在初始化DeepSeek查询管理器...")
            query_manager = AIQueryManager()
        else:
            print("正在初始化OLLAMA查询管理器...")
            query_manager = OllamaQueryManager()
            
        print(query_manager._generate_table_description())

        # 示例问题列表
        questions = [
            # 租户相关查询
            "查询所有租户的基本信息",
            "显示每个租户的用户数量",
            "查找最近一个月新增的租户",
            "统计每个租户的合同总数",
            
            # 用户相关查询
            # "查询所有系统用户信息",
            # "显示每个租户下的用户列表",
            # "查找最近一周内创建的用户",
            # "统计每个租户的活跃用户数量",
            
            # 合同相关查询
            # "查询所有有效的合同信息",
            # "显示最近签订的合同列表",
            # "查找即将到期的合同",
            # "统计每个租户的合同总金额",
            
            # 项目相关查询
            # "查询所有进行中的项目",
            # "显示每个项目的预算使用情况",
            # "查找进度超过80%的项目",
            # "统计每个租户的项目数量",
            
            # # 收付款相关查询
            # "查询最近一个月的收款记录",
            # "显示每个合同的付款计划",
            # "查找逾期未付款的合同",
            # "统计每个租户的收款总额",
            
            # # 发票相关查询
            # "查询最近开具的发票",
            # "显示每个合同的发票开具情况",
            # "查找未开具发票的收款记录",
            # "统计每个租户的发票总金额",
            
            # # 综合查询
            # "查询某个租户的合同、项目和发票汇总信息",
            # "显示项目预算超支的合同详情",
            # "查找发票金额与合同金额不匹配的记录",
            # "统计每个租户的合同执行情况"
        ]
        
        # 处理每个问题
        for i, question in enumerate(questions, 1):
            print(f"\n处理第 {i}/{len(questions)} 个问题...")
            try:
                print(f"\n问题: {question}")
                result = query_manager.query(question)
                print_result(question, result, output_dir)
            except Exception as e:
                error_result = {
                    'success': False,
                    'error': str(e),
                    'traceback': traceback.format_exc()
                }
                print_result(question, error_result, output_dir)
                
    except Exception as e:
        print(f"程序执行出错: {str(e)}")
        print("\n错误详情:")
        print(traceback.format_exc())

if __name__ == "__main__":
    main() 